[Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat - #9
Closed
JaredforReal wants to merge 51 commits into
Closed
JaredforReal wants to merge 51 commits into
JaredforReal wants to merge 51 commits into
Conversation
…oject#55375) Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
…3835) Signed-off-by: WikAlie <194917782+wei-core@users.noreply.github.com> Co-authored-by: WikAlie <194917782+wei-core@users.noreply.github.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Signed-off-by: linitra24 <renshuang.zhou@daocloud.io> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
vllm-project#55448) Signed-off-by: Luca Motz <luca.motz@icloud.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
…lm-project#49410) Signed-off-by: zhejiangxiaomai <zhenhui.zhao@intel.com>
…GB PCIe (vllm-project#55511) Signed-off-by: Baki Burak Ogun <63836730+bakiburakogun@users.noreply.github.com>
…t#51898) Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com>
…hing and spec-decoding (vllm-project#53614) Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: levius <2114377220@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
There was a problem hiding this comment.
🟢 Approval recommended
The changes are narrowly scoped, consistent with existing call patterns, and avoid extra allocation/copy without altering tensor shapes for the NoPE (pe_dim=0) case.
Pull request overview
Enables faster SM100 prefill for GLM-5.3-Flash’s NoPE MLA layout by (1) allowing masked-MHA for the (64, 512, 256, 0, 256) dimension tuple and (2) avoiding an unnecessary K concatenation allocation/copy when the RoPE portion is empty.
Changes:
- Extend SM100 masked-MHA allow-list to include GLM-5.3-Flash’s NoPE (qk_nope=256, qk_rope=0, v=256) layout.
- Optimize MLA K construction by returning
k_nopedirectly whenk_pehas a zero last-dimension (no RoPE), skipping allocation/copy.
File summaries
| File | Description |
|---|---|
| vllm/model_executor/layers/attention/sparse_mla_attention.py | Adds GLM-5.3-Flash NoPE dimension tuple to SM100 masked-MHA availability allow-list. |
| vllm/model_executor/layers/attention/mla_attention.py | Skips K concatenation when k_pe is empty to avoid unnecessary allocation/copy in NoPE MLA. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…ct#55075) Signed-off-by: Alex <jihui.huang@daocloud.io> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
…T2 (vllm-project#55461) Signed-off-by: Wang Zupeng <zupenwang@gmail.com> Co-authored-by: Codex <noreply@openai.com>
vllm-project#54362) Signed-off-by: Will Hampson <whamp@ggl.slmail.me> Co-authored-by: Will Hampson <whamp@ggl.slmail.me> Co-authored-by: Or Ozeri <oro@il.ibm.com>
vllm-project#54288) Signed-off-by: almogtavor <almogtavor@gmail.com> Signed-off-by: Almog Tavor <70065337+almogtavor@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com>
…ct#55529) Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Codex <noreply@openai.com>
…-project#54659) Signed-off-by: zhouyou9505 <zhouyou9505@gmail.com> Co-authored-by: Cursor Grok 4.6 <noreply@cursor.com>
…-project#55457) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Codex <noreply@openai.com>
…project#55455) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Signed-off-by: Wenhua Cheng <wenhua.cheng@intel.com>
…uted experts (vllm-project#53161) Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Cursor Agent <agent@cursor.com>
…llm-project#54975) Signed-off-by: Big2Wheel <1404207590@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
…ect#55642) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
…rEngine (vllm-project#41567) Signed-off-by: Teng Ma <sima.mt@alibaba-inc.com> Signed-off-by: Tianyu Guo <guoty@inferact.ai> Signed-off-by: Zhou ziheng <jiaranran2@gmail.com> Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: Zhou ziheng <jiaranran2@gmail.com> Co-authored-by: jiangkuaixue123 <jiangxiaozhou111@163.com> Co-authored-by: OpenAI Codex <codex@openai.com>
…oject#55653) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
…llm-project#54643) Signed-off-by: zhewenl <zhewenl@users.noreply.github.com> Co-authored-by: zhewenl <zhewenl@users.noreply.github.com>
…r BF16 variant (vllm-project#53586) Signed-off-by: Zihua Wu <zihuaw@nvidia.com> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
…ternal_tokens (vllm-project#47505) Signed-off-by: Alex <alex.tech.lab@outlook.com> Signed-off-by: Alex-ai-future <Alex-ai-future@users.noreply.github.com>
…latforms (vllm-project#55660) Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
…tention prefill backend GLM-5.3-Flash (qk_nope 256, qk_rope 0, v 256) was not in FlashAttnPrefillBackend.supports_mla_dimensions, so sparse MLA fell back to "No MLA prefill backend supports this model" and every prefill token went through the per-token top-k MQA kernel. The kernels for qk_head_dim 256 / v_head_dim 256 are the ones already used for the (192, 64, 256) layout. Same 7 lines as in the SM90 wiring PR (vllm-project#55385); kept as a separate commit so it can be dropped once that lands. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
…4, 512, 256, 0, 256) layout wire_sm90_kernels registers the (256, 0, 256) MLA dimensions with the FlashAttention prefill backend, which gives GLM-5.3-Flash the dense-MHA prefill path for sequences up to index_topk (2048). Sequences above that still went through the per-token top-k MQA kernel on SM100, which is KV-gather bound (5.9 ms/layer per 16k-token chunk, ~4.8 ms of it pure HBM traffic), because the masked-MHA allow-list only knew the (128, 512, 128, 64, 128) and (64, 512, 192, 64, 256) layouts. The masked kernel path is identical for qk_head_dim 256 / v 256, so add the GLM layout; the existing FLASHINFER_MLA_SPARSE TP4/TP8 thresholds (masked MHA up to 36k / 64k tokens) apply. Prefill TTFT on 4x GB300 (TP4, together with the FlashKDA and decode cleanup PRs): 8x2048 462 -> 363 ms, 2x8192 382 -> 317 ms, 32k/256 c16 TTFT 3876 -> 3425 ms. First use of a new sequence-length bucket JIT compiles the FA4 mask_mod variant (~7 s); the FA4 MLA prefill warmup does not cover mask_mod yet. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
With qk_rope_head_dim == 0 the concatenated K is just k_nope; return it instead of allocating and copying a same-sized tensor (134 MB per layer for a 16k-token chunk on GLM-5.3-Flash). Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
…project#54404) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
…rror (vllm-project#50257) Signed-off-by: shaolila <shaolila@buaa.edu.cn> Co-authored-by: shaolila <shaolila@buaa.edu.cn>
…ect#53580) Signed-off-by: Marceli Fylcek <marceli.fylcek@intel.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Signed-off-by: Sean Westfall <sean.westfall@gmail.com>
JaredforReal
force-pushed
the
perf/glm53-flash-masked-mha
branch
from
September 7, 2026 10:39
5299590 to
85928b9
Compare
…54890) Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
…LE spec decode (vllm-project#52771) Signed-off-by: Kam Basra <kameldipbasra@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
…-project#52651) Signed-off-by: Artur Fierka <artur.fierka@intel.com>
…50195) Signed-off-by: Francisco Javier Arceo <farceo@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
This was referenced Sep 7, 2026
Owner
Author
|
Upstream draft: vllm-project#55738 |
Signed-off-by: Jared Wen <w13431838023@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
One of three independent GLM-5.3-Flash perf PRs (#7: decode cleanups; #8: FlashKDA prefill; this one: masked-MHA prefill). Each is benchmarked on its own below. Rebased onto main (previously on
wire_sm90_kernels); the one piece it needed from there is carried as its own first commit.Summary
On main, GLM-5.3-Flash's MLA layout (qk_nope 256, qk_rope 0, v 256) is unknown to the FlashAttention prefill backend and to the masked-MHA allow-list, so
MLACommonImpllogsNo MLA prefill backend supports this modeland every prefill token goes through the per-token top-k MQA kernel. Profiling shows that kernel is KV-gather bound at the HBM roofline: 5.9 ms/layer per 16k-token chunk, and still 1.6 ms/layer for 8x2048 where dense MHA needs ~0.2 ms.Three commits:
FlashAttnPrefillBackend(7 lines). Same lines as in the SM90 wiring PR [perf] wire FA and FlashMLA for sm90 GLM5Next NoPE SparseMLA vllm-project/vllm#55385; kept separate so it can be dropped once that lands. This alone gives the dense-MHA prefill path for sequences ≤ index_topk (2048)._is_masked_mha_availableonly knew (128, 512, 128, 64, 128) and (64, 512, 192, 64, 256); GLM-5.3-Flash is (64, 512, 256, 0, 256). The kernels are identical for qk_head_dim 256 / v 256, so sequences between 2k and the existing FLASHINFER_MLA_SPARSE thresholds (36k at TP4, 64k at TP8) now use FA4 masked MHA instead of the per-token top-k MQA kernel._concat_k_nope_k_pereturnsk_nopedirectly when there is no RoPE part instead of allocating and copying a same-sized tensor (134 MB/layer per 16k chunk).Performance of this PR alone
Measured on top of #7 + #8 (main + #7 + #8) vs the same + this PR, 4x GB300, TP4,
--attention-backend FLASHINFER_MLA_SPARSE --max-model-len 69632 --max-num-seqs 256 --max-num-batched-tokens 16384, prefix caching disabled (hit rate 0.0% checked in the server log),vllm bench serverandom dataset with warmups, back-to-back in the same session. Decode tok/s is the steady-state window value; TPOT is the per-request median.Prefill up to the masked-MHA threshold (≤ 36k at TP4): −8.5% TTFT for 8x2048 (dense MHA), −7.4% for 8x32k (masked MHA), −4.4% TTFT in the 32k-context serving point. 2x8192 shows −4.6%, but that 16-request point has a ~4% run-to-run noise floor (last column of the ablation below), so treat it as indicative only. 64k prompts are above the threshold and stay on the sparse kernel; the −3.5% there (0.0% noise floor) is the K-concat removal alone. Decode is unchanged (±0.8%).
Ablation of the whole series
Cumulative, same session, same settings: main → +#7 → +#8 → +#9 (this PR) → main again.
Caveat: runtime JIT
The first request in each new sequence-length bucket JIT-compiles the FA4 CuTeDSL
mask_modvariant (BlackwellFusedMultiHeadAttentionForward) plus the Triton_scatter_topk_single_req_kernel(~7 s each;jit_monitorwarns).FA4MLAPrefillKernel.get_warmup_keysdoes not cover mask_mod, so this is a pre-existing gap for every masked-MHA model; a warmup for the masked variants should be a follow-up before relying on this in production.Accuracy (per build, same session)
Each rung of the ablation was also evaluated on its own: gsm8k (1319 questions, 5-shot, greedy, lm_eval
local-completions), prompt-logprob agreement on real 4k/12k/30k/60k prompts (mean |Δ logprob| per token vs the main build; the "main again" row is the run-to-run noise floor of this FP8 model), and a needle-in-a-haystack retrieval set (12 codes per length at 6k/16k/30k tokens, greedy,reasoning_effort=low), which exercises the 2k-36k prefill range where #9 switches kernels.| build | gsm8k flexible-extract % | gsm8k strict-match % | prompt-logprob mean|Δ| vs main (4k/12k/30k/60k) | next token | needle hits (6k/16k/30k) | needle, 3 more passes (concurrency 6) |
|---|---|---|---|---|---|---|
| main | 93.33 ± 0.69 | 93.10 ± 0.70 | (reference) | | 12/12 / 12/12 / 12/12 | 34/36 / 36/36 / 35/36 |
| +#7 | 93.03 ± 0.70 | 92.95 ± 0.71 | 0.040 / 0.055 / 0.051 / 0.108 | all same | 12/12 / 12/12 / 11/12 | 35/36 / 36/36 / 34/36 |
| +#7 +#8 | 93.25 ± 0.69 | 93.18 ± 0.69 | 0.099 / 0.072 / 0.060 / 0.136 | all same | 12/12 / 11/12 / 11/12 | - |
| +#7 +#8 +#9 | 93.10 ± 0.70 | 93.10 ± 0.70 | 0.099 / 0.072 / 0.060 / 0.127 | all same | 12/12 / 12/12 / 11/12 | 35/36 / 34/36 / 36/36 |
| main again | 92.87 ± 0.71 | 92.87 ± 0.71 | 0.041 / 0.055 / 0.051 / 0.118 | all same | 12/12 / 12/12 / 12/12 | - |
+#7vs main equals the main-vs-main noise floor to three digits, i.e. [Perf][GLM-5.3-Flash] Decode hot-path cleanups: strided KDA recurrent inputs, NoPE MQA query without concat, no duplicate router GEMM #7 is numerically equivalent to main (expected: the strided KDA kernel is bit-identical, the bmm into a transposedoutwas verified identical, and the router logits actually used were always the ones MoERunner recomputes). [Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) #8 and [Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat #9 raise the 4k |Δ| to ~2.4x the noise floor (FlashKDA and FA4 accumulate in a different order than the Triton chunk / per-token sparse kernels); 12k-60k stay at the noise floor and the next token is identical everywhere.Duplicate-work check
gh pr list --repo vllm-project/vllm --state open --search "NoPE masked MHA 256"/"GLM-5.3-Flash prefill": vllm-project#55385 covers the FA dims (commit 1 here is the same 7 lines); vllm-project#54951 / vllm-project#55222 / vllm-project#55543 touch the indexer prefill sharding, the indexer workspace and SM90 fp8 KV. None adds the masked-MHA layout or the K-concat skip.Tests
tests/kernels/test_mhc_kernels.py,tests/kernels/test_glm5next_kda_recurrent_strided.pyunaffected; prompt-logprob A/B on real 4k/12k/30k/60k prompts vs main: mean|Δ| 0.06-0.12 vs a main-vs-main noise floor of 0.04-0.11 (≈2.5x at 4k where dense/masked MHA replaces the sparse kernel), identical next tokens; gsm8k 1319 questions (5-shot, greedy): main 93.03% ± 0.70, full series 92.65% ± 0.72 (measured on the three PRs together).pre-commit run ruff-check / ruff-format: passed.~/notes/glm53flash-perf/REPORT.md, ablation logs in~/notes/glm53flash-perf/logs/e2e/ablate_*.